import os
import fire
import sys
import time

# Add the project root to the Python path
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..'))
sys.path.append(project_root)

from omics_bert_label_encoding.label.read_label_embeddings import LabelEmbeddingReader

def main(dataset_path: str = os.path.join(project_root, "embeddings_local/ukbiobank/v4-labels_embeddings")):
    """Browse embeddings in a dataset."""
    with LabelEmbeddingReader(dataset_path) as reader:
        print(f"Browsing embeddings in {dataset_path}")

        # 统计embedding有多少个
        print(f"Total embeddings: {reader.count_embeddings()}")

        # 示例：取出key为'p3_i0'对应的embedding
        embedding_p48_i0 = reader.get_label_embedding('p48_i0')
        print(f"Embedding for key 'p48_i0': {embedding_p48_i0}")

        # 示例：取出['p3_i0', 'pzbtb17_i2']对应的embedding
        embeddings_batch = reader.get_labels_embeddings_batch(['p3_i0', 'p48_i0'])
        print(f"Embeddings for keys ['p3_i0', 'p48_i0']: {embeddings_batch}")

        # 打印前10个embedding
        print("Printing first 10 embeddings...")
        
        for i, (key, embedding) in enumerate(reader.iterate_embeddings()):
            if i >= 10:
                print("\nReached 10 embeddings, stopping.")
                break
                
            print(f"\n--- Embedding {i+1} ---")
            print(f"Key: {key}")
            print(f"Embedding: {embedding}")


if __name__ == '__main__':
    fire.Fire(main)